"""
在 Neo4j 中构建医疗知识图谱
选材自开源项目(刘焕勇，中国科学院软件研究所)，数据集来自互联网爬虫数据
@Time: 2021/5/10 10:50
@Author:
@File: build_medical_graph.py
"""
import os
import json
from py2neo import Graph, Node
from typing import Set, List, Tuple
from pydantic import BaseModel, Field


class GraphNodeStore(BaseModel):
    """
    用于存储所有节点数据的数据对象
    """
    drugs: List = Field(default_factory=list,
                        title='药品')
    foods: List = Field(default_factory=list,
                        title='食物')
    checks: List = Field(default_factory=list,
                         title='检查')
    departments: List = Field(default_factory=list,
                              title='科室')
    producers: List = Field(default_factory=list,
                            title='药品大类')
    diseases: List = Field(default_factory=list,
                           title='疾病')
    symptoms: List = Field(default_factory=list,
                           title='症状')


class Disease(BaseModel):
    """
    用于描述一个疾病及其属性的数据对象
    """
    name: str = Field(default_factory=str)
    desc: str = Field(default_factory=str)
    prevent: str = Field(default_factory=str)
    cause: str = Field(default_factory=str)
    easy_get: str = Field(default_factory=str)
    cure_department: str = Field(default_factory=str)
    cure_way: str = Field(default_factory=str)
    cure_prob: str = Field(default_factory=str)
    cure_lasttime: str = Field(default_factory=str)
    symptom: str = Field(default_factory=str)
    cured_prob: str = Field(default_factory=str)
    get_prob: str = Field(default_factory=str)


class GraphRelStore(BaseModel):
    """
    用于存储所有实体关系的数据对象
    """
    department: List = Field(default_factory=list,
                             title='科室－科室关系')
    not_eat: List = Field(default_factory=list,
                          title='疾病－忌吃食物关系')
    do_eat: List = Field(default_factory=list,
                         title='疾病－宜吃食物关系')
    recommend_eat: List = Field(default_factory=list,
                                title='疾病－推荐吃食物关系')
    common_drug: List = Field(default_factory=list,
                              title='疾病－通用药品关系')
    recommend_drug: List = Field(default_factory=list,
                                 title='疾病－热门药品关系')
    check: List = Field(default_factory=list,
                        title='疾病－检查关系')
    drug_producer: List = Field(default_factory=list,
                                title='厂商－药物关系')
    symptom: List = Field(default_factory=list,
                          title='疾病症状关系')
    accompany: List = Field(default_factory=list,
                            title='疾病并发关系')
    category: List = Field(default_factory=list,
                           title='疾病与科室之间的关系')


class MedicalGraph:
    def __init__(self):
        cur_dir = '/'.join(os.path.abspath(__file__).split('/')[:-1])
        self.data_path = os.path.join(cur_dir, 'resource/medical2.json')
        self.g = Graph("http://localhost:7474", auth=("neo4j", "password"))

    def read_nodes(self) -> Tuple[GraphNodeStore, List[Disease], GraphRelStore]:
        """
        读取文件数据，并将读取的数据接入返回的三个对象中
        :return:
        """
        node_store = GraphNodeStore()  # 所有节点
        disease_infos: List[Disease] = []  # 疾病信息
        rel_store = GraphRelStore()  # 节点实体的关系

        count = 0
        for data in open(self.data_path, encoding='utf8'):
            disease = Disease()
            count += 1
            print(count)
            data_json = json.loads(data)
            disease_name = data_json['name']
            disease.name = disease_name
            node_store.diseases.append(disease_name)

            if 'symptom' in data_json:
                node_store.symptoms += data_json['symptom']
                for symptom in data_json['symptom']:
                    rel_store.symptom.append([disease_name, symptom])

            if 'acompany' in data_json:
                for acompany in data_json['acompany']:
                    rel_store.accompany.append([disease_name, acompany])

            if 'desc' in data_json:
                disease.desc = data_json['desc']

            if 'prevent' in data_json:
                disease.prevent = data_json['prevent']

            if 'cause' in data_json:
                disease.cause = data_json['cause']

            if 'get_prob' in data_json:
                disease.get_prob = data_json['get_prob']

            if 'easy_get' in data_json:
                disease.easy_get = data_json['easy_get']

            if 'cure_department' in data_json:
                cure_department = data_json['cure_department']
                if len(cure_department) == 1:
                    rel_store.category.append([disease_name, cure_department[0]])
                if len(cure_department) == 2:
                    big = cure_department[0]
                    small = cure_department[1]
                    rel_store.department.append([small, big])
                    rel_store.category.append([disease_name, small])

                disease.cure_department = cure_department
                rel_store.department += cure_department

            if 'cure_way' in data_json:
                disease.cure_way = data_json['cure_way']

            if 'cure_lasttime' in data_json:
                disease.cure_lasttime = data_json['cure_lasttime']

            if 'cured_prob' in data_json:
                disease.cure_prob = data_json['cured_prob']

            if 'common_drug' in data_json:
                common_drug = data_json['common_drug']
                for drug in common_drug:
                    rel_store.common_drug.append([disease_name, drug])
                node_store.drugs += common_drug

            if 'recommand_drug' in data_json:
                recommand_drug = data_json['recommand_drug']
                node_store.drugs += recommand_drug
                for drug in recommand_drug:
                    rel_store.recommend_drug.append([disease_name, drug])

            if 'not_eat' in data_json:
                not_eat = data_json['not_eat']
                for _not in not_eat:
                    rel_store.not_eat.append([disease_name, _not])

                node_store.foods += not_eat
                do_eat = data_json['do_eat']
                for _do in do_eat:
                    rel_store.do_eat.append([disease_name, _do])

                node_store.foods += do_eat
                recommand_eat = data_json['recommand_eat']

                for _recommand in recommand_eat:
                    rel_store.recommend_eat.append([disease_name, _recommand])
                node_store.foods += recommand_eat

            if 'check' in data_json:
                check = data_json['check']
                for _check in check:
                    rel_store.check.append([disease_name, _check])
                node_store.checks += check
            if 'drug_detail' in data_json:
                drug_detail = data_json['drug_detail']
                producer = [i.split('(')[0] for i in drug_detail]
                rel_store.drug_producer += [[i.split('(')[0], i.split('(')[-1].replace(')', '')] for i in drug_detail]
                node_store.producers += producer
            disease_infos.append(disease)
        return node_store, disease_infos, rel_store

    def create_node(self, label: str, nodes: Set[str]) -> None:
        """
        创建节点
        :param label: 节点的标签名
        :param nodes: 由所有节点名称组成的集合
        """
        count = 0
        for node_name in nodes:
            node = Node(label, name=node_name)
            self.g.create(node)
            count += 1
            print(count, len(nodes))
        return

    def create_diseases_nodes(self, disease_list: List[Disease]) -> None:
        """
        创建知识图谱中心疾病的节点
        :param disease_list: 由所有表示疾病信息的对象组成的列表
        """
        count = 0
        for disease in disease_list:
            node = Node("Disease", **disease.dict())
            self.g.create(node)
            count += 1
            print(count)
        return

    def create_graph_nodes_and_rels(self):
        node_store, disease_list, rel_store = self.read_nodes()
        self._create_graph_nodes(node_store, disease_list)
        self._create_graph_rels(rel_store)

    def _create_graph_nodes(self, node_store: GraphNodeStore, disease_list: List[Disease]) -> None:
        """
        创建知识图谱实体节点类型schema
        :return:
        """
        self.create_diseases_nodes(disease_list)
        self.create_node('Drug', set(node_store.drugs))
        self.create_node('Food', set(node_store.foods))
        self.create_node('Check', set(node_store.checks))
        self.create_node('Department', set(node_store.departments))
        self.create_node('Producer', set(node_store.producers))
        self.create_node('Symptom', set(node_store.symptoms))

    def _create_graph_rels(self, rel_store: GraphRelStore) -> None:
        """
        创建实体关系边
        """
        self._create_relationship('Disease', 'Food', rel_store.recommend_eat, 'recommand_eat', '推荐食谱')
        self._create_relationship('Disease', 'Food', rel_store.not_eat, 'no_eat', '忌吃')
        self._create_relationship('Disease', 'Food', rel_store.do_eat, 'do_eat', '宜吃')
        self._create_relationship('Department', 'Department', rel_store.department, 'belongs_to', '属于')
        self._create_relationship('Disease', 'Drug', rel_store.common_drug, 'common_drug', '常用药品')
        self._create_relationship('Producer', 'Drug', rel_store.drug_producer, 'drugs_of', '生产药品')
        self._create_relationship('Disease', 'Drug', rel_store.recommend_drug, 'recommand_drug', '好评药品')
        self._create_relationship('Disease', 'Check', rel_store.check, 'need_check', '诊断检查')
        self._create_relationship('Disease', 'Symptom', rel_store.symptom, 'has_symptom', '症状')
        self._create_relationship('Disease', 'Disease', rel_store.accompany, 'acompany_with', '并发症')
        self._create_relationship('Disease', 'Department', rel_store.category, 'belongs_to', '所属科室')

    def _create_relationship(self,
                             start_node_type: str,
                             end_node_type: str,
                             edges: List[List[str]],
                             rel_type: str,
                             rel_name: str):
        """
        创建实体关联边
        :param start_node_type: 起点的点类型
        :param end_node_type: 终点的点类型
        :param edges: 由边组成的列表，每个元素是由起点名称和终点名称构成的列表
        :param rel_type: 关系的类型
        :param rel_name: 关系的名称
        :return:
        """
        count = 0
        # 去重处理
        set_edges = []
        for edge in edges:
            set_edges.append('###'.join(edge))
        all_edges = len(set(set_edges))
        for edge in set(set_edges):
            edge = edge.split('###')
            p = edge[0]
            q = edge[1]
            query = "match(p:%s),(q:%s) where p.name='%s'and q.name='%s' create (p)-[rel:%s{name:'%s'}]->(q)" % (
                start_node_type, end_node_type, p, q, rel_type, rel_name)
            try:
                self.g.run(query)
                count += 1
                print(rel_type, count, all_edges)
            except Exception as e:
                print(e)
        return

    def export_data(self):
        """导出数据"""
        node_store, disease_list, rel_store = self.read_nodes()
        f_drug = open('drug.txt', 'w+')
        f_food = open('food.txt', 'w+')
        f_check = open('check.txt', 'w+')
        f_department = open('department.txt', 'w+')
        f_producer = open('producer.txt', 'w+')
        f_symptom = open('symptoms.txt', 'w+')
        f_disease = open('disease.txt', 'w+')

        f_drug.write('\n'.join(list(node_store.drugs)))
        f_food.write('\n'.join(list(node_store.foods)))
        f_check.write('\n'.join(list(node_store.checks)))
        f_department.write('\n'.join(list(node_store.departments)))
        f_producer.write('\n'.join(list(node_store.producers)))
        f_symptom.write('\n'.join(list(node_store.symptoms)))
        f_disease.write('\n'.join(list(node_store.diseases)))

        f_drug.close()
        f_food.close()
        f_check.close()
        f_department.close()
        f_producer.close()
        f_symptom.close()
        f_disease.close()

        return


if __name__ == '__main__':
    handler = MedicalGraph()
    # handler.export_data()
    handler.create_graph_nodes_and_rels()
